package org.wenshu.ai.example;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

import static dev.langchain4j.data.message.UserMessage.userMessage;
import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;

public class _05_Memory {

    public static void main(String[] args) throws ExecutionException, InterruptedException {

        OpenAiStreamingChatModel model = OpenAiStreamingChatModel.builder()
            .baseUrl("https://maas-api.cn-huabei-1.xf-yun.com/v1")
            .apiKey("your-key")
            .modelName("xdeepseekv3")
                .build();

        ChatMemory chatMemory = TokenWindowChatMemory.withMaxTokens(1000, new HuggingFaceTokenizer());

        SystemMessage systemMessage = SystemMessage.from(
                "你是一位高级开发者，正在向另一位高级开发者解释，"
                        + "你正在开发的项目是一个使用Java后端、" +
                        "Oracle数据库和Spring Data JPA的电子商务平台");
        chatMemory.add(systemMessage);


        UserMessage userMessage1 = userMessage(
                "如何优化大规模电子商务平台的数据库查询？"
                        + "请用三到五行简短回答。");
        chatMemory.add(userMessage1);

        System.out.println("[用户]: " + userMessage1.text());
        System.out.print("[AI]: ");

        CompletableFuture<AiMessage> futureAiMessage = new CompletableFuture<>();

        StreamingChatResponseHandler handler = new StreamingChatResponseHandler() {

            @Override
            public void onPartialResponse(String partialResponse) {
                System.out.print(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                futureAiMessage.complete(completeResponse.aiMessage());
            }

            @Override
            public void onError(Throwable throwable) {
            }
        };

        model.chat(chatMemory.messages(), handler);
        chatMemory.add(futureAiMessage.get());

        UserMessage userMessage2 = userMessage(
                "能给出第一点的具体实现示例吗？" +
                        "请简短回答，最多10行代码。");
        chatMemory.add(userMessage2);

        System.out.println("\n\n[用户]: " + userMessage2.text());
        System.out.print("[AI]: ");

        model.chat(chatMemory.messages(), handler);
    }
}
